package com.mm.cloth.rmbg;

import ai.onnxruntime.*;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.util.Log;

import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.Collections;

/**
 * 基于ONNX Runtime的Watermark-Removal推理助手
 * 专门用于水印检测和去除
 */
public class WatermarkRemovalInferenceHelper {
    private static final String TAG = "WatermarkRemovalInference";
    
    private OrtSession session;
    private OrtEnvironment env;
    
    // 模型参数
    private static final String INPUT_NAME = "input";
    private static final String MASK_INPUT_NAME = "mask"; // 可选的遮罩输入
    private static final String OUTPUT_NAME = "output";
    
    // 图像标准化参数（通常用于Watermark-Removal模型）
    private static final float[] MEAN = {0.485f, 0.456f, 0.406f};
    private static final float[] STD = {0.229f, 0.224f, 0.225f};
    
    // 模型输入尺寸（根据实际模型调整）
    private static final int MODEL_INPUT_SIZE = 512;
    
    /**
     * 初始化ONNX Runtime环境
     * @param assets AssetManager
     * @param modelPath 模型文件路径
     */
    public void initModel(AssetManager assets, String modelPath) throws Exception {
        Log.d(TAG, "初始化Watermark-Removal ONNX Runtime环境，模型路径: " + modelPath);
        
        try {
            // 检查ONNX Runtime是否可用
            try {
                Class.forName("ai.onnxruntime.OrtEnvironment");
                Log.d(TAG, "ONNX Runtime类可用");
            } catch (ClassNotFoundException e) {
                throw new Exception("ONNX Runtime不可用", e);
            }
            
            // 检查系统架构兼容性
            String arch = System.getProperty("os.arch");
            Log.d(TAG, "系统架构: " + arch);
            
            // 创建ONNX Runtime环境（使用更安全的方式）
            try {
                env = OrtEnvironment.getEnvironment();
                Log.d(TAG, "ONNX Runtime环境创建成功");
            } catch (Exception e) {
                Log.e(TAG, "ONNX Runtime环境创建失败", e);
                throw new Exception("ONNX Runtime环境初始化失败: " + e.getMessage(), e);
            }
            
            // 配置会话选项（使用最保守的设置）
            OrtSession.SessionOptions options = new OrtSession.SessionOptions();
            try {
                options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.NO_OPT); // 禁用优化避免兼容性问题
                options.setIntraOpNumThreads(1); // 使用单线程避免并发问题
                options.setInterOpNumThreads(1); // 限制操作间线程
                Log.d(TAG, "会话选项配置完成");
            } catch (Exception e) {
                Log.e(TAG, "会话选项配置失败", e);
                throw new Exception("会话选项配置失败: " + e.getMessage(), e);
            }
            
            // 检查模型文件
            InputStream modelStream = assets.open(modelPath);
            int available = modelStream.available();
            Log.d(TAG, "模型文件大小: " + (available / (1024*1024)) + " MB");
            
            if (available == 0) {
                modelStream.close();
                throw new Exception("模型文件为空");
            }
            
            // 加载模型数据
            byte[] modelBytes = new byte[available];
            int totalRead = 0;
            while (totalRead < available) {
                int read = modelStream.read(modelBytes, totalRead, available - totalRead);
                if (read == -1) break;
                totalRead += read;
            }
            modelStream.close();
            
            Log.d(TAG, "模型数据加载完成: " + totalRead + " bytes");
            
            // 验证模型数据
            if (totalRead < 1000) {
                throw new Exception("模型文件太小，可能损坏");
            }
            
            // 创建会话（添加更多错误检查）
            try {
                Log.d(TAG, "开始创建ONNX会话...");
                session = env.createSession(modelBytes, options);
                Log.d(TAG, "ONNX会话创建成功");
                
                // 验证会话是否可用
                if (session == null) {
                    throw new Exception("会话创建返回null");
                }
                
            } catch (Exception e) {
                Log.e(TAG, "ONNX会话创建失败", e);
                throw new Exception("ONNX会话创建失败: " + e.getMessage(), e);
            }
            
            Log.d(TAG, "Watermark-Removal ONNX Runtime模型加载成功");
            
            // 打印模型信息
            printModelInfo();
            
        } catch (Exception e) {
            Log.e(TAG, "Watermark-Removal ONNX Runtime模型初始化失败", e);
            // 清理资源
            if (session != null) {
                try { session.close(); } catch (Exception ignored) {}
                session = null;
            }
            if (env != null) {
                try { env.close(); } catch (Exception ignored) {}
                env = null;
            }
            throw e;
        }
    }
    
    /**
     * 打印模型信息
     */
    private void printModelInfo() {
        try {
            Log.d(TAG, "=== Watermark-Removal模型信息 ===");
            Log.d(TAG, "输入节点数量: " + session.getInputNames().size());
            Log.d(TAG, "输出节点数量: " + session.getOutputNames().size());
            
            for (String inputName : session.getInputNames()) {
                Log.d(TAG, "输入节点: " + inputName);
            }
            
            for (String outputName : session.getOutputNames()) {
                Log.d(TAG, "输出节点: " + outputName);
            }
            
        } catch (Exception e) {
            Log.w(TAG, "获取模型信息失败", e);
        }
    }
    
    /**
     * 使用遮罩去除水印
     * @param originalImage 原始图像
     * @param maskImage 遮罩图像（白色区域为水印位置）
     * @return 去水印后的图像
     */
    public Bitmap removeWatermarkWithMask(Bitmap originalImage, Bitmap maskImage) throws Exception {
        if (session == null || env == null) {
            throw new IllegalStateException("模型未初始化");
        }
        
        if (originalImage == null || maskImage == null) {
            throw new IllegalArgumentException("输入图像或遮罩不能为空");
        }
        
        Log.d(TAG, "开始带遮罩的水印去除推理");
        Log.d(TAG, "原始图像尺寸: " + originalImage.getWidth() + "x" + originalImage.getHeight());
        Log.d(TAG, "遮罩图像尺寸: " + maskImage.getWidth() + "x" + maskImage.getHeight());
        
        long startTime = System.currentTimeMillis();
        OnnxTensor imageTensor = null;
        OnnxTensor maskTensor = null;
        
        try {
            // 1. 预处理图像
            float[] imageData = preprocessImage(originalImage);
            long[] imageShape = {1, 3, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE}; // NCHW格式
            
            // 2. 预处理遮罩
            float[] maskData = preprocessMask(maskImage);
            long[] maskShape = {1, 1, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE}; // NCHW格式
            
            Log.d(TAG, "预处理完成，图像数据长度: " + imageData.length + ", 遮罩数据长度: " + maskData.length);
            
            // 3. 构建输入张量
            try {
                Log.d(TAG, "创建输入张量...");
                imageTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(imageData), imageShape);
                maskTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(maskData), maskShape);
                Log.d(TAG, "输入张量创建成功");
            } catch (Exception e) {
                Log.e(TAG, "输入张量创建失败", e);
                throw new Exception("输入张量创建失败: " + e.getMessage(), e);
            }
            
            // 4. 执行推理
            long inferenceStart = System.currentTimeMillis();
            Log.d(TAG, "开始ONNX推理...");
            
            // 构建输入映射
            java.util.Map<String, OnnxTensor> inputs = new java.util.HashMap<>();
            inputs.put(INPUT_NAME, imageTensor);
            inputs.put(MASK_INPUT_NAME, maskTensor);
            
            try (OrtSession.Result result = session.run(inputs)) {
                long inferenceEnd = System.currentTimeMillis();
                Log.d(TAG, "ONNX推理完成，耗时: " + (inferenceEnd - inferenceStart) + "ms");
                
                // 5. 获取输出结果
                try {
                    Log.d(TAG, "获取推理结果...");
                    Object outputObject = result.get(0).getValue();
                    if (outputObject == null) {
                        throw new Exception("推理结果为null");
                    }
                    
                    Log.d(TAG, "转换输出数据...");
                    float[] outputData = convertOutputToFloatArray(outputObject);
                    if (outputData == null || outputData.length == 0) {
                        throw new Exception("输出数据为空");
                    }
                    
                    // 6. 后处理：将输出转换为Bitmap
                    Log.d(TAG, "开始后处理...");
                    Bitmap resultBitmap = postprocessOutput(outputData, originalImage.getWidth(), originalImage.getHeight());
                    if (resultBitmap == null) {
                        throw new Exception("后处理失败，结果为null");
                    }
                    
                    long totalTime = System.currentTimeMillis() - startTime;
                    Log.d(TAG, "Watermark-Removal推理完成，总耗时: " + totalTime + "ms");
                    
                    return resultBitmap;
                    
                } catch (Exception e) {
                    Log.e(TAG, "输出处理失败", e);
                    throw new Exception("输出处理失败: " + e.getMessage(), e);
                }
            }
            
        } catch (OutOfMemoryError e) {
            Log.e(TAG, "Watermark-Removal推理内存不足", e);
            System.gc();
            throw new Exception("内存不足，请尝试使用较小的图片", e);
        } catch (Exception e) {
            Log.e(TAG, "Watermark-Removal推理失败", e);
            throw e;
        } finally {
            // 确保释放张量资源
            if (imageTensor != null) {
                try {
                    imageTensor.close();
                } catch (Exception e) {
                    Log.w(TAG, "释放图像张量资源失败", e);
                }
            }
            if (maskTensor != null) {
                try {
                    maskTensor.close();
                } catch (Exception e) {
                    Log.w(TAG, "释放遮罩张量资源失败", e);
                }
            }
        }
    }
    
    /**
     * 自动检测并去除水印
     * @param originalImage 原始图像
     * @return 去水印后的图像
     */
    public Bitmap removeWatermarkAuto(Bitmap originalImage) throws Exception {
        if (session == null || env == null) {
            throw new IllegalStateException("模型未初始化");
        }
        
        if (originalImage == null) {
            throw new IllegalArgumentException("输入图像不能为空");
        }
        
        Log.d(TAG, "开始自动水印检测和去除推理");
        Log.d(TAG, "原始图像尺寸: " + originalImage.getWidth() + "x" + originalImage.getHeight());
        
        long startTime = System.currentTimeMillis();
        OnnxTensor tensor = null;
        
        try {
            // 1. 预处理图像
            float[] inputData = preprocessImage(originalImage);
            long[] shape = {1, 3, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE}; // NCHW格式
            
            Log.d(TAG, "预处理完成，输入数据长度: " + inputData.length);
            
            // 2. 构建输入张量
            try {
                Log.d(TAG, "创建输入张量...");
                tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), shape);
                Log.d(TAG, "输入张量创建成功");
            } catch (Exception e) {
                Log.e(TAG, "输入张量创建失败", e);
                throw new Exception("输入张量创建失败: " + e.getMessage(), e);
            }
            
            // 3. 执行推理
            long inferenceStart = System.currentTimeMillis();
            Log.d(TAG, "开始ONNX推理...");
            try (OrtSession.Result result = session.run(Collections.singletonMap(INPUT_NAME, tensor))) {
                long inferenceEnd = System.currentTimeMillis();
                Log.d(TAG, "ONNX推理完成，耗时: " + (inferenceEnd - inferenceStart) + "ms");
                
                // 4. 获取输出结果
                try {
                    Log.d(TAG, "获取推理结果...");
                    Object outputObject = result.get(0).getValue();
                    if (outputObject == null) {
                        throw new Exception("推理结果为null");
                    }
                    
                    Log.d(TAG, "转换输出数据...");
                    float[] outputData = convertOutputToFloatArray(outputObject);
                    if (outputData == null || outputData.length == 0) {
                        throw new Exception("输出数据为空");
                    }
                    
                    // 5. 后处理：将输出转换为Bitmap
                    Log.d(TAG, "开始后处理...");
                    Bitmap resultBitmap = postprocessOutput(outputData, originalImage.getWidth(), originalImage.getHeight());
                    if (resultBitmap == null) {
                        throw new Exception("后处理失败，结果为null");
                    }
                    
                    long totalTime = System.currentTimeMillis() - startTime;
                    Log.d(TAG, "自动Watermark-Removal推理完成，总耗时: " + totalTime + "ms");
                    
                    return resultBitmap;
                    
                } catch (Exception e) {
                    Log.e(TAG, "输出处理失败", e);
                    throw new Exception("输出处理失败: " + e.getMessage(), e);
                }
            }
            
        } catch (OutOfMemoryError e) {
            Log.e(TAG, "Watermark-Removal推理内存不足", e);
            System.gc();
            throw new Exception("内存不足，请尝试使用较小的图片", e);
        } catch (Exception e) {
            Log.e(TAG, "Watermark-Removal推理失败", e);
            throw e;
        } finally {
            // 确保释放张量资源
            if (tensor != null) {
                try {
                    tensor.close();
                } catch (Exception e) {
                    Log.w(TAG, "释放张量资源失败", e);
                }
            }
        }
    }

    /**
     * 预处理图像
     * @param image 输入图像
     * @return 预处理后的float数组
     */
    private float[] preprocessImage(Bitmap image) {
        // 调整图像尺寸到模型输入尺寸
        Bitmap resizedImage = Bitmap.createScaledBitmap(image, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, true);

        // 获取像素数据
        int[] pixels = new int[MODEL_INPUT_SIZE * MODEL_INPUT_SIZE];
        resizedImage.getPixels(pixels, 0, MODEL_INPUT_SIZE, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);

        // 转换为NCHW格式的float数组
        int totalPixels = MODEL_INPUT_SIZE * MODEL_INPUT_SIZE;
        float[] inputData = new float[3 * totalPixels];

        // 分别处理R、G、B三个通道
        for (int c = 0; c < 3; c++) {
            for (int h = 0; h < MODEL_INPUT_SIZE; h++) {
                for (int w = 0; w < MODEL_INPUT_SIZE; w++) {
                    int pixelIndex = h * MODEL_INPUT_SIZE + w;
                    int pixel = pixels[pixelIndex];

                    float value;
                    if (c == 0) { // R通道
                        value = Color.red(pixel) / 255.0f;
                    } else if (c == 1) { // G通道
                        value = Color.green(pixel) / 255.0f;
                    } else { // B通道
                        value = Color.blue(pixel) / 255.0f;
                    }

                    // ImageNet标准化: (value - mean) / std
                    value = (value - MEAN[c]) / STD[c];

                    // NCHW格式存储
                    int outputIndex = c * totalPixels + pixelIndex;
                    inputData[outputIndex] = value;
                }
            }
        }

        // 清理临时bitmap
        if (resizedImage != image) {
            resizedImage.recycle();
        }

        return inputData;
    }

    /**
     * 预处理遮罩
     * @param mask 输入遮罩
     * @return 预处理后的float数组
     */
    private float[] preprocessMask(Bitmap mask) {
        // 调整遮罩尺寸到模型输入尺寸
        Bitmap resizedMask = Bitmap.createScaledBitmap(mask, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, true);

        // 获取像素数据
        int[] pixels = new int[MODEL_INPUT_SIZE * MODEL_INPUT_SIZE];
        resizedMask.getPixels(pixels, 0, MODEL_INPUT_SIZE, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);

        // 转换为单通道float数组
        int totalPixels = MODEL_INPUT_SIZE * MODEL_INPUT_SIZE;
        float[] maskData = new float[totalPixels];

        for (int i = 0; i < totalPixels; i++) {
            int pixel = pixels[i];
            // 将遮罩转换为0-1值（白色=1，黑色=0）
            float value = Color.red(pixel) / 255.0f;
            maskData[i] = value;
        }

        // 清理临时bitmap
        if (resizedMask != mask) {
            resizedMask.recycle();
        }

        return maskData;
    }

    /**
     * 转换ONNX输出为float数组
     */
    private float[] convertOutputToFloatArray(Object outputObject) {
        try {
            if (outputObject instanceof float[][][][]) {
                // 4维张量格式 [batch, channel, height, width]
                float[][][][] output4D = (float[][][][]) outputObject;

                int batch = output4D.length;
                int channels = output4D[0].length;
                int height = output4D[0][0].length;
                int width = output4D[0][0][0].length;

                Log.d(TAG, "输出张量形状: [" + batch + ", " + channels + ", " + height + ", " + width + "]");

                // 转换为1维数组
                float[] result = new float[channels * height * width];
                int index = 0;
                for (int c = 0; c < channels; c++) {
                    for (int h = 0; h < height; h++) {
                        for (int w = 0; w < width; w++) {
                            result[index++] = output4D[0][c][h][w];
                        }
                    }
                }

                return result;

            } else if (outputObject instanceof float[]) {
                return (float[]) outputObject;
            } else {
                throw new RuntimeException("不支持的输出格式: " + outputObject.getClass().getName());
            }

        } catch (Exception e) {
            Log.e(TAG, "输出转换失败", e);
            throw new RuntimeException("输出转换失败", e);
        }
    }

    /**
     * 后处理输出数据
     * @param outputData 模型输出数据
     * @param targetWidth 目标图像宽度
     * @param targetHeight 目标图像高度
     * @return 处理后的Bitmap
     */
    private Bitmap postprocessOutput(float[] outputData, int targetWidth, int targetHeight) {
        // 首先创建模型输出尺寸的bitmap
        Bitmap modelOutputBitmap = Bitmap.createBitmap(MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, Bitmap.Config.ARGB_8888);
        int[] pixels = new int[MODEL_INPUT_SIZE * MODEL_INPUT_SIZE];

        int totalPixels = MODEL_INPUT_SIZE * MODEL_INPUT_SIZE;

        // 检查输出数据长度
        int expectedLength = 3 * totalPixels; // RGB 3通道
        if (outputData.length != expectedLength) {
            Log.w(TAG, "输出数据长度不匹配，期望: " + expectedLength + ", 实际: " + outputData.length);
        }

        for (int h = 0; h < MODEL_INPUT_SIZE; h++) {
            for (int w = 0; w < MODEL_INPUT_SIZE; w++) {
                int pixelIndex = h * MODEL_INPUT_SIZE + w;

                // 从NCHW格式提取RGB值
                float r = outputData[pixelIndex]; // R通道
                float g = outputData[totalPixels + pixelIndex]; // G通道
                float b = outputData[2 * totalPixels + pixelIndex]; // B通道

                // 反标准化: value = value * std + mean
                r = r * STD[0] + MEAN[0];
                g = g * STD[1] + MEAN[1];
                b = b * STD[2] + MEAN[2];

                // 确保值在[0,1]范围内
                r = Math.max(0.0f, Math.min(1.0f, r));
                g = Math.max(0.0f, Math.min(1.0f, g));
                b = Math.max(0.0f, Math.min(1.0f, b));

                // 转换为0-255范围
                int red = (int) (r * 255);
                int green = (int) (g * 255);
                int blue = (int) (b * 255);

                pixels[pixelIndex] = Color.argb(255, red, green, blue);
            }
        }

        modelOutputBitmap.setPixels(pixels, 0, MODEL_INPUT_SIZE, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);

        // 调整到目标尺寸
        Bitmap resultBitmap;
        if (targetWidth != MODEL_INPUT_SIZE || targetHeight != MODEL_INPUT_SIZE) {
            resultBitmap = Bitmap.createScaledBitmap(modelOutputBitmap, targetWidth, targetHeight, true);
            modelOutputBitmap.recycle();
        } else {
            resultBitmap = modelOutputBitmap;
        }

        Log.d(TAG, "后处理完成，生成图像: " + resultBitmap.getWidth() + "x" + resultBitmap.getHeight());
        return resultBitmap;
    }

    /**
     * 释放资源
     */
    public void release() {
        try {
            if (session != null) {
                session.close();
                session = null;
            }
            if (env != null) {
                env.close();
                env = null;
            }
            Log.d(TAG, "Watermark-Removal ONNX Runtime资源已释放");
        } catch (Exception e) {
            Log.e(TAG, "释放Watermark-Removal ONNX Runtime资源失败", e);
        }
    }

    /**
     * 检查是否已初始化
     */
    public boolean isInitialized() {
        return session != null && env != null;
    }
}
